import os
import shutil
import wave
from glob import glob

from config import Config
from models import WakeupModule


def check_file(file):
    wakeup = WakeupModule(Config["weight_path"]).eval()
    with wave.open(file, 'r') as wf:
        wav_data = wf.readframes(wf.getnframes())

    outputs = wakeup.recognize(wav_data)
    print(file, outputs)


def check_folder(root, tag, moveto=None, threshold=0.1):
    print(root, tag)
    if moveto and not os.path.exists(moveto):
        os.mkdir(moveto)

    wakeup = WakeupModule(Config["weight_path"]).eval()
    files = glob(root + '/*.wav')
    for file in files:
        with wave.open(file, 'r') as wf:
            wav_data = wf.readframes(wf.getnframes())
        outputs = wakeup.recognize(wav_data)

        if abs(outputs[0] - tag) >= threshold:
            print(file, outputs[0])
            if moveto:
                name = os.path.basename(file)
                des = os.path.join(moveto, name)
                shutil.move(file, des)


def pickup_folder(root, tag, moveto=None, threshold=0.1):
    print(root, tag)
    if moveto and not os.path.exists(moveto):
        os.mkdir(moveto)

    wakeup = WakeupModule(Config["weight_path"]).eval()
    files = glob(root + '/*.wav')
    for file in files:
        with wave.open(file, 'r') as wf:
            wav_data = wf.readframes(wf.getnframes())
        outputs = wakeup.recognize(wav_data)

        if abs(outputs[0] - tag) < threshold:
            print(file, outputs[0])
            if moveto:
                name = os.path.basename(file)
                des = os.path.join(moveto, name)
                shutil.move(file, des)


def move_files(frm, to):
    files = glob(frm + '/*.wav')
    for f in files:
        des = os.path.join(to, os.path.basename(f))
        shutil.move(f, des)


if __name__ == '__main__':
    def check_repeat():
        subs = glob("../data/train/*")
        for sub in subs:
            check_folder(sub + "/-/repeat", 0,     sub + "/-/err", 0.4)
            check_folder(sub + "/+/repeat", 1,     sub + "/+/err", 0.4)
            check_folder(sub + "/±/repeat", 0.5,   sub + "/±/err", 0.3)

    def check_all():
        subs = glob("../data/train/*")
        for sub in subs:
            check_folder(sub + "/-", 0, sub + "/-/error", 0.4)
            check_folder(sub + "/+", 1, sub + "/+/error", 0.4)
            check_folder(sub + "/±", 0.5, sub + "/±/error", 0.3)

    def move_err_to_repeat():
        subs = glob("../data/train/*")
        for sub in subs:
            move_files(sub + "/-/err", sub + "/-/repeat")
            move_files(sub + "/+/err", sub + "/+/repeat")
            move_files(sub + "/±/err", sub + "/±/repeat")

    def pickup_TBD():
        subs = glob("../data/train/*")
        for sub in subs:
            pickup_folder(sub + "/TBD", 0,    sub + "/TBD/-", 0.1)
            pickup_folder(sub + "/TBD", 1,    sub + "/TBD/+", 0.1)
            pickup_folder(sub + "/TBD", 0.5,  sub + "/TBD/±", 0.05)

    # move_err_to_repeat()
    # check_repeat()
    check_all()
    # pickup_TBD()
    # pickup_folder("../data/train/record/-/err", 1, "../data/temp", 0.4)
    # check_folder("../data/train/20210607/-", 0, "../data/train/20210607/-/err", 0.1)

    print('Finished!')
